//
// KeyedLists.cxx
//
// Created by TB on 11/17/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <cmath>
#include <iostream>
#include <limits>

#include <boost/foreach.hpp>

#include <TRandom.h>

#include "tppmc/KeyedLists.h"

//#define mNQBins        50
//#define mQMin    2.0
//#define mQMax  102.0

namespace tppmc {
   
   // \remark
   // Couldn't all this be done more easily with a TH1, using
   // its Interpolate() method?
   
   KeyedLists::KeyedLists(int nQBins,
                                      double QMin,
                                      double QMax)

//   : mQBins(nQBins, QMin, QMax)
   : mNQBins(nQBins)
   , mQMin(QMin)
   , mQMax(QMax) {
//      std::cout << "KeyedLists::KeyedLists("<<nQBins<<
//      ","<<QMin<<","<<QMax<<")"<<std::endl;
   }
   
   
   bool KeyedLists::Insert(int parton) {
      return mTable.insert(std::make_pair(parton,
                                          Array(mNQBins, 0.))).second;
   }
   
   
   double KeyedLists::At(int parton, unsigned QBin) const {
      if(not Exists(parton, QBin)) {
         return std::numeric_limits<double>::quiet_NaN();
      } // if
      Table::const_iterator i = mTable.find(parton);
      return i->second.at(QBin);
   }
   
   
   double KeyedLists::operator()(int parton, unsigned QBin) const {
      return At(parton, QBin);
   }
   
   
   bool KeyedLists::Exists(int parton, unsigned QBin) const {
      Table::const_iterator i = mTable.find(parton);
      if(i not_eq mTable.end()) {
         return QBin < mNQBins;
      } // if
      return false;
   }
   
   
   void KeyedLists::Set(int parton, unsigned QBin, double value) {
      if(Exists(parton, QBin)) {
         Table::iterator i = mTable.find(parton);
         i->second.at(QBin) = value;
      } // if
   }
   
   
   double KeyedLists::Interpolate(int parton, double Q) const {
      if(not Exists(parton, 0)) {
         return 0.;
      }
      
//      double bin_size = (mQMax - mQMin) / (float)mNQBins; 
      // Find the bin number containing the requested Q value.
      int sidx_low = static_cast<int>((Q - mQMin) / BinSize());
//      int sidx_low = mQBins.FindBin(Q);
      
      if(sidx_low < 0) {
         sidx_low = 0;
      } // if
      else if(unsigned(sidx_low) >= (mNQBins - 1)) {
         sidx_low = mNQBins - 2;
      } // else if
      
      // Bin number above the requested Q value.
      int sidx_high = sidx_low + 1; 
      
      double IpPiplus =
      At(parton, sidx_low) +
      ((Q - (mQMin + (sidx_low * BinSize()))) / BinSize()) *
      (At(parton, sidx_high) - At(parton, sidx_low));
      
      return IpPiplus;
   }
   
   //
   // class FragmentationTable
   //
   
   FragmentationTable::~FragmentationTable() { }
   
   /**
    Initialise the table with uniform binning in Q.
    Set nBins Q bins, each of width (max-min)/nBins.
    @param parton  Identifier for the table
    @param nBins   Number of Q-squared bins
    @param min     Minimum Q-squared
    @param max     Maximum Q-squared
    @param hadrons Optional list of hadron species
    */
   FragmentationTable::FragmentationTable(//int parton,
                                          int nBins,
                                          double min,
                                          double max,
                                          const std::vector<int>& hadrons)
   : mPrototype("", "", nBins, min, max) {
      mPrototype.SetDirectory(NULL);
      BOOST_FOREACH(int hadron, hadrons) {
         Insert(hadron);
      } // BOOST_FOREACH
   }
   
   /**
    \overload
    Initialise the table with non-uniform binning in Q.
    @param parton  Identifier for the table
    @param edges A list of Q bin edges. The number of resultant bins is
    edges.size() - 1. Each entry in edges must be larger than the last.
    */
   FragmentationTable::FragmentationTable(//int parton,
                                          std::vector<double> edges,
                                          const std::vector<int>& hadrons)
   : mPrototype("", "", edges.size() - 1, &(edges.at(0))) {
      mPrototype.SetDirectory(NULL);
      BOOST_FOREACH(int hadron, hadrons) {
         Insert(hadron);
      } // BOOST_FOREACH
   }
   
   /**
    Normalise the current table contents such that, in each Q bin,
    the sum over all hadrons equals norm.
    */
   void FragmentationTable::Normalise(double norm) {
      
      // Accumulate sum over species
      Array sum(mPrototype);
      sum.SetDirectory(NULL);
      
      BOOST_FOREACH(const Entry& entry, mTable) {
         for(int i(1); i <= entry.second.GetNbinsX(); ++i) {
            sum.Fill(sum.GetBinCenter(i), entry.second.GetBinContent(i));
         } // for
      } // BOOST_FOREACH
      
      // Divide contents, Q bin by Q bin
      BOOST_FOREACH(Entry& entry, mTable) {
         entry.second.Divide(&sum);
         entry.second.Scale(norm);
      } // BOOST_FOREACH
   }
   
   /**
    Returns the value for species in the bin containing Q.
    */
   double FragmentationTable::At(int species, double Q) const {
      if(not Contains(species, Q)) {
         return std::numeric_limits<double>::quiet_NaN();
      } // if
      Table::const_iterator i = mTable.find(species);
      int bin = mPrototype.FindBin(Q);
      return i->second.GetBinContent(bin);
   }


   /** \overload */
   //      const double& At(int species, double Q) const;
   /**
    Set the value for species in the bin containing Q.
    */
   bool FragmentationTable::Set(int species, double Q, double value) {
      if(Contains(species, Q)) {
         int bin = mPrototype.FindBin(Q);
         mTable[species].SetBinContent(bin, value);
         return true;
      } // if
      return false;
   }


   /**
    Returns true if a value is stored for species and the bin
    containing Q.
    */
   bool FragmentationTable::Contains(int species, double Q) const {
      
      // Is there an entry for this species?
      Table::const_iterator entry = mTable.find(species);
      if(mTable.end() == entry) {
         return false; // Nope, sorry
      } // if
      
      // We have this species, what about Q?
      // TH1 bins are numbered [1, N].
      int bin = mPrototype.FindBin(Q);
      return bin >= 1 and bin <= mPrototype.GetNbinsX();
   }


   /**
    Expand the table with an entry for species.
    Binning in Q is the same as for all other entries.
    Initialise each Q bin with the value initial.
    Returns true if the new entry is inserted, false if not.
    */
   bool FragmentationTable::Insert(int species, double initial) {
      
      // Is there an entry for this species?
      // If so do nothing and return false.
      Table::const_iterator entry = mTable.find(species);
      if(mTable.end() not_eq entry) {
         return false;
      } // if
      
      // No entry currently, so copy the prototype and initialise
      mTable.insert(std::make_pair(species, Array(mPrototype)));
      mTable[species].SetDirectory(NULL);
      for(int i(1); i <= mTable[species].GetNbinsX(); ++i) {
         mTable[species].SetBinContent(i, initial);
      } // for
      
      return true;
   }


   void FragmentationTable::Print(Option_t* /* not used */) const {
      
      std::ios::fmtflags flags = std::cout.flags(); // Store initial flags
      
      std::cout.precision(4);
      std::cout.width(10);
      
      // Start with the Q bin edges
      std::cout << "Bin edges:";
      for(int i(1); i <= mPrototype.GetNbinsX() + 1; ++i) {
         std::cout << mPrototype.GetBinLowEdge(i) << "\t";
      } // for
      std::cout << std::endl;
      
      // Accumulate sum over species so we can print a total
      Array sum(mPrototype);
      sum.SetDirectory(NULL);
      
      // Now the contents of each Q bin for each species
      BOOST_FOREACH(const Entry& entry, mTable) {
         std::cout << entry.first << '\t';
         for(int i(1); i <= entry.second.GetNbinsX(); ++i) {
            std::cout << entry.second.GetBinContent(i) << '\t';
            sum.Fill(sum.GetBinCenter(i), entry.second.GetBinContent(i));
         } // for
         std::cout << std::endl;
      } // BOOST_FOREACH
      
      std::cout << "sum\t";
      for(int i(1); i <= sum.GetNbinsX(); ++i) {
         std::cout << sum.GetBinContent(i) << '\t';
      } // for
      std::cout << std::endl;
      
      std::cout.flags(flags); // Restore initial flags
   }


   std::pair<int, bool> FragmentationTable::Random(double Q) {
      const double value = gRandom->Uniform(0., 1.);
      double sum(0.);
      // Find the fragmentation probability at this exact Q (via interpolation
      // between bins). Add up the values until we exceed the random number.
      // Then that's our random hadron.
      BOOST_FOREACH(Entry& entry, mTable) {
         sum += entry.second.Interpolate(Q);
         if(sum > value) {
            return std::make_pair(entry.first, true);
         } // if
      } // BOOST_FOREACH
      return std::make_pair(std::numeric_limits<int>::min(), false);
   }


   int FragmentationTable::Bin(double Q) const {
      int bin = mPrototype.FindBin(Q);
      // Histogram bin lies in range [1, N] so subtract 1 for [0, N).
      if(bin >= 1 and bin <= mPrototype.GetNbinsX()) return bin - 1;
      return -1;
   }

} // namespace tppmc
